#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File:pic_to_file.py
@Author:Josco
@Date:2022/10/08 11:09:35
@Version:1.0
@Description: 将jpg文件解码
'''
import cv2 as cv
import numpy as np
from hash_check import *
import time
import multiprocessing
import os.path
from PIL import Image
from io import BytesIO
from hash_check import *
# import cupy as cp



def scan_dir(dir_path, pic_type):
    """扫描目录"""
    file_list = []
    for top, dirs, nondirs in os.walk(dir_path):
        for item in nondirs:
            if item.split('.')[-1] == pic_type:
                file_list.append(os.path.join(top, item))
    return file_list
# 把保存bmp文件的目录下所有的bmp文件转为jpg文件
def bmp_to_jpg(file_path):
    file_list = scan_dir(file_path, "bmp")
    for one_file in file_list:
        bmp = Image.open(one_file)
        output_buffer = BytesIO()
        bmp.save(output_buffer,format ='JPEG', quality=90)
        byte_data = output_buffer.getvalue()
        a = os.path.splitext(one_file)[0]
        a = a.split("./data2/")[1]
        tmp_url = './data3/'+a +'.jpg'
        try:
            with open(tmp_url, 'wb') as f:
                f.write(byte_data)
        except Exception as e:
            pass

def color_to_num(mat):
    """
        使用此函数把矩阵中的0-255的值映射到0-3空间中
    """
    mat = np.where((mat>=0) & (mat < 43), 0, mat)
    mat = np.where((mat>=43) & (mat < 128), 1, mat)
    mat = np.where((mat>=128) & (mat < 213), 2, mat)
    mat = np.where((mat>=213) & (mat <= 255), 3, mat)
    return mat

def num_to_strbit(mat):
    """
        使用此函数把矩阵中0-3的数字转化为01字节流字符串
    """
    mat = np.where(mat=="0", "00", mat)
    mat = np.where(mat=="1", "01", mat)
    mat = np.where(mat=="2", "10", mat)
    mat = np.where(mat=="3", "11", mat)
    return mat

def pic_to_file(pic_name):
    """采用像素中心4个像素的均值当作色块颜色"""
    # 加载jpg图像
    n = 5 # 单张伪二维码放大倍数
    top_bottom = 3 # 单张伪二维码上下黑边
    left_right = 3  # 单张伪二维码左右黑边
    image = cv.imread(pic_name)
    image = np.array(image)
    
    # 去除黑边
    image = image[top_bottom*n:-(top_bottom*n),left_right*n:-(left_right*n),:]

    b = image[:,:,0] # 得到蓝色通道
    g = image[:,:,1] # 得到绿色通道
    r = image[:,:,2] # 得到红色通道
    b = b.astype("uint")
    g = g.astype("uint")
    r = r.astype("uint")
    
    # r1 = (r[int(n/2)::n, int(n/2)::n] + r[int(n/2)-1::n, int(n/2)-1::n]) / 2
    r1 = r[int(n/2)::n, int(n/2)::n]
    r1 = color_to_num(r1)
    # g1 = (g[int(n/2)::n, int(n/2)::n] + g[int(n/2)-1::n, int(n/2)-1::n]) / 2
    g1 = g[int(n/2)::n, int(n/2)::n]
    g1 = color_to_num(g1)
    # b1 = (b[int(n/2)::n, int(n/2)::n] + b[int(n/2)-1::n, int(n/2)-1::n]) / 2
    b1 = b[int(n/2)::n, int(n/2)::n]
    # print(b1.astype("uint8"))
    b1 = color_to_num(b1)

    # # 按rgb通道连接数组
    img = np.dstack((b1, g1, r1))
    img = img.astype("uint8")
    
    h,w,c = img.shape
    img = img.reshape(h,-1) # 转为h行,w*c列的二维数组
    
    # 处理文件原始数据，并更具原始数据计算md5校验码
    check_img = img[2:] # 获取第2行到最后一行的文件原始数据（03矩阵），用于md5校验
    check_info = img[:2].astype(str)
    # 对解码后的check_img进行md5校验, 计算md5的值
    do_hash_check = hash_check()
    decode_md5_check_info = do_hash_check.md5(str(check_img))
    # 把文件信息转为[0 1.... 1 0]的一维数组
    check_img = check_img.astype("str") # 把文件numpy数组转为numpy字符串
    check_img = num_to_strbit(check_img) # 把0-3映射为00,01,10,11
    check_img = check_img.reshape(-1) # 转为一维数组
    check_img_data = "".join(check_img) # 数组中每个元素连接为str字符串
    
    check_img_data = np.fromstring(check_img_data, np.int8) - 48
    # check_img_data = np.array(list(check_img_data), dtype="int") # 从字符串直接转为np.array数组
    

    # 获取前两列的基本信息和校验信息
    # 将前两列03区间映射为00，01，10，11
    check_info = num_to_strbit(check_info)
    check_info = check_info.reshape(-1)
    check_info_data = "".join(check_info) # 获取前两行基本信息和校验信息，存储前两行01bit流数据的str

    # 按jiezhidaowo这个尾部校验值切分前两行数据, 并留下完整的基本信息和校验信息
    check_info_data = check_info_data.split("00110001001100100011001100110100001100110011001000110001")
    check_info_data = check_info_data[0]
    # 调用hash_check函数进行校验
    check_info_data = do_hash_check.bitarray2str(check_info_data) # 从bit流转为字节流
    check_info_data = check_info_data.split(" ") # 按空格切割校验信息
    code_md5_check_info = check_info_data[-2] #取得md5校验信息

    fin_check_info_data = [check_info_data[0], check_info_data[1]+" "+check_info_data[2],check_info_data[3]+" "+check_info_data[4],check_info_data[5]+" "+check_info_data[6],check_info_data[7], check_info_data[8], check_info_data[9], check_info_data[10],check_info_data[11], int(check_info_data[12]),check_info_data[13]]
    # print(check_img_data)
    # print(fin_check_info_data)

    # 对前两行的校验信息和经过md5后的校验信息做对比，若一致，则返回check_info 及 原始文件01数组
    if decode_md5_check_info == code_md5_check_info:
        # print(pic_name + "校验成功-----")
        return [fin_check_info_data, check_img_data]

    print(pic_name + "伪二维码解码出错，编码的校验码为：", code_md5_check_info)
    print("解码的校验码为：", decode_md5_check_info)
    return False

def deal_return_file_info(all_return_info):
    """对解码返回的文件信息和基本信息进行拼接"""

    remain_data = int(all_return_info[0][0][7])
    # print(remain_data)

    reversed_data = sorted(all_return_info, key=lambda x:x[0][9])   # 将数组按照第几章来排序

    all_reversed_data = reversed_data[0][1]
    for reversed_data_item in reversed_data[1:]:
        all_reversed_data = np.concatenate((all_reversed_data, reversed_data_item[1]))
    all_reversed_data = all_reversed_data[:-remain_data]

    bytes_all_reversed_data = np.packbits(all_reversed_data)
    bytes_all_reversed_data.tofile("./demo111.pdf")


def main():
    # # 把保存在bmp_file_path中的bmp文件转为jpg文件并存储到./data3/下面
    # bmp_file_path = "./data2/"
    # bmp_to_jpg(bmp_file_path)

    time_start = time.time()
    # 获取jpg_file_path中的jpg的张数
    jpg_file_path = "./data3/"
    file_list = scan_dir(jpg_file_path, "jpg")
    all_return_info = []

    # 采用map_asyncD的方式
    pool = multiprocessing.Pool(processes = 10)
    r = pool.map_async(pic_to_file, file_list, callback = all_return_info.extend)
    pool.close()
    pool.join()   # 调用join之前，先调用close函数，否则会出错。执行完close后不会有新的进程加入到pool,join函数等待所有子进程结束
    print("Sub-process(es) done.")

    deal_return_file_info(all_return_info)
    time_end = time.time()
    time_sum = time_end - time_start
    print(time_sum)

if __name__ == '__main__':
    main()